import collections
import glob
import json
import os
import os.path as osp
import uuid
import xml.etree.ElementTree as ET

import cv2
import labelme
import numpy as np
import pycocotools.mask
import tqdm
from PIL import Image


def write_xml(imgname,filepath,labeldicts):                     #参数imagename是图片名（无后缀）
    root = ET.Element('Annotation')                             #创建Annotation根节点
    ET.SubElement(root, 'filename').text = str(imgname)         #创建filename子节点（无后缀）
    sizes = ET.SubElement(root,'size')                          #创建size子节点            
    ET.SubElement(sizes, 'width').text = str(img.shape[1])                #没带脑子直接写了原图片的尺寸......
    ET.SubElement(sizes, 'height').text = str(img.shape[0])  
    ET.SubElement(sizes, 'depth').text = str(img.shape[2])                    #图片的通道数：img.shape[2]
    for labeldict in labeldicts:
        objects = ET.SubElement(root, 'object')                 #创建object子节点
        ET.SubElement(objects, 'name').text = labeldict['name']        #BDD100K_10.names文件中  
                                                                       #的类别名
        ET.SubElement(objects, 'pose').text = 'Unspecified'
        ET.SubElement(objects, 'truncated').text = '0'
        ET.SubElement(objects, 'difficult').text = '0'
        bndbox = ET.SubElement(objects,'bndbox')
        ET.SubElement(bndbox, 'xmin').text = str(int(labeldict['xmin']))
        ET.SubElement(bndbox, 'ymin').text = str(int(labeldict['ymin']))
        ET.SubElement(bndbox, 'xmax').text = str(int(labeldict['xmax']))
        ET.SubElement(bndbox, 'ymax').text = str(int(labeldict['ymax']))
    tree = ET.ElementTree(root)
    tree.write(filepath, encoding='utf-8')



annotations_path = r'/media/hjh/workdir/0_Deep_Learning/datasets/ciwa_beimian_VOC2007/val_xml/'   #生成的xml文件需要保存的路径
file_dir = r'/media/hjh/workdir/0_Deep_Learning/datasets/ciwa_beimian_VOC2007/val/'    #json文件路径
json_path = os.listdir(file_dir)

label_files = glob.glob(osp.join(file_dir, '*.json'))
i = 0
for filename in tqdm.tqdm(label_files):
    if filename.endswith('.json'):
        # root = open(os.path.join(file_dir, fp), 'rb')
        # data = json.load(root)
        label_file = labelme.LabelFile(filename=filename)
        img_name = osp.splitext(osp.basename(filename))[0]
        img = labelme.utils.img_data_to_arr(label_file.imageData)
        
        masks = {}                                     # for area
        segmentations = collections.defaultdict(list)  # for segmentation
        for shape in label_file.shapes:
            points = shape['points']
            label = shape['label']
            group_id = shape.get('group_id')
            shape_type = shape.get('shape_type', 'polygon')
            mask = labelme.utils.shape_to_mask(
                img.shape[:2], points, shape_type
            )

            if group_id is None:
                group_id = uuid.uuid1()

            instance = (label, group_id)

            if instance in masks:
                masks[instance] = masks[instance] | mask
            else:
                masks[instance] = mask

        labeldicts = []
        for instance, mask in masks.items():
            mask = np.asfortranarray(mask.astype(np.uint8))
            mask = pycocotools.mask.encode(mask)
            # area = float(pycocotools.mask.area(mask))
            bbox = pycocotools.mask.toBbox(mask).flatten().tolist()
            new_dict = {'name': instance[0],    # class_name
                        'difficult': '0',
                        'xmin': bbox[0],                      
                        'ymin': bbox[1],
                        'xmax': bbox[0] + bbox[2],
                        'ymax': bbox[1] + bbox[3]
                        }
            labeldicts.append(new_dict)
        write_xml(img_name + '.jpg', annotations_path + img_name + '.xml', labeldicts)
        # print(str(img_name + '.xml')+' 转换成功')
        i = i + 1
print('共',i,'次')




